# -*- coding: utf-8 -*-

from __future__ import annotations

import os
# 用于显示进度条
from tqdm import tqdm
# 用于发起网络请求
import requests
# 用于多线程操作
import multitasking
import signal
# 导入 retry 库以方便进行下载出错重试
from retry import retry
from typing import List
from bs4 import BeautifulSoup as Soup

signal.signal(signal.SIGINT, multitasking.killall)

# 请求头
headers = {
    'User-Agent': 'Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/102.0.0.0 Safari/537.36'
}
# 定义 1 MB 多少为 B
MB = 1024**2


class InnerMindWithHeart:
    host: str
    lecture_category: str
    audio_download_path: str = ''
    audio_source_list: List[str]

    def __init__(self):
        self.host = 'https://yishineihua.com'
        self.lecture_category = f'{self.host}/archives/category'
        self.audio_source_list = []

    def collect_article_audios(self):
        lecture_res = requests.get(self.url)
        lecture_soup = Soup(lecture_res.text, 'html.parser')
        lecture_content = lecture_soup.find(id='content')

        print('start to collect lecture articles ...')
        for item in lecture_content.find_all('article'):
            article_url = item.find('h1').find('a').attrs['href']

            article_res = requests.get(article_url)
            article_soup = Soup(article_res.text, 'html.parser')
            audio = article_soup.find('audio')
            if audio is not None:
                audio_source = audio.text
                self.audio_source_list.append(audio_source)
                print(f'finished parsing article audio: {audio_source}')
            else:
                print(f'article audio not found, article url: {article_url}')

    def download_article_audios(self):
        if len(self.audio_source_list) > 0:
            for url_item in self.audio_source_list:
                file_name = os.path.basename(url_item)
                self.download_mp3(url_item, file_name)

    @staticmethod
    def _split(end: int, step: int) -> list[tuple[int, int]]:
        # 分多块
        parts = [(start, min(start + step, end))
                 for start in range(0, end, step)]
        return parts

    @staticmethod
    def _get_file_size(url: str, raise_error: bool = False) -> int:
        """
        获取文件大小

        Parameters
        ----------
        url : 文件直链
        raise_error : 如果无法获取文件大小，是否引发错误

        Return
        ------
        文件大小（B为单位）
        如果不支持则会报错
        """
        response = requests.head(url)
        file_size = response.headers.get('Content-Length')
        if file_size is None:
            if raise_error is True:
                raise ValueError('该文件不支持多线程分段下载！')
            return file_size
        return int(file_size)

    def download_mp3(self, url: str, file_name: str, retry_times: int = 3, each_size=16 * MB) -> None:
        """
        根据文件直链和文件名下载文件，参考文档：https://zhuanlan.zhihu.com/p/369531344

        Parameters
        ----------
        url : 文件下载链接
        file_name : 文件名
        retry_times: 可选的，每次连接失败重试次数
        Return
        ------
        None

        """
        file_path = os.path.join(self.audio_download_path, file_name)
        if os.path.exists(file_path):
            print(f'{file_name} already exists, pass')
            return

        f = open(file_path, 'wb')
        file_size = self._get_file_size(url)

        @retry(tries=retry_times)
        @multitasking.task
        def start_download(start_idx: int, end_idx: int) -> None:
            """
            根据文件起止位置下载文件

            Parameters
            ----------
            start_idx : 开始位置
            end_idx : 结束位置
            """
            _headers = headers.copy()
            # 分段下载的核心
            _headers['Range'] = f'bytes={start_idx}-{end_idx}'
            # 发起请求并获取响应（流式）
            response = session.get(url, headers=_headers, stream=True)
            # 每次读取的流式响应大小
            chunk_size = 128
            # 暂存已获取的响应，后续循环写入
            chunks = []
            for chunk in response.iter_content(chunk_size=chunk_size):
                # 暂存获取的响应
                chunks.append(chunk)
                # 更新进度条
                bar.update(chunk_size)
            f.seek(start_idx)
            for chunk in chunks:
                f.write(chunk)
            # 释放已写入的资源
            del chunks

        session = requests.Session()
        # 分块文件如果比文件大，就取文件大小为分块大小
        each_size = min(each_size, file_size)

        # 分块
        parts = self._split(file_size, each_size)
        print(f'分块数：{len(parts)}')
        # 创建进度条
        bar = tqdm(total=file_size, desc=f'下载文件：{file_name}')
        for part in parts:
            start, end = part
            start_download(start, end)
        # 等待全部线程结束
        multitasking.wait_for_tasks()
        f.close()
        bar.close()
